#pragma once
#include <iostream>
#include <utility>
#include <cassert>

using namespace std;

template<class K, class V>
struct AVLTreeNode
{
    pair<K, V> _kv;
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;
	int _bf;//平衡因子

	AVLTreeNode(const pair<K, V>& kv)
		: _kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _bf(0)
	{}
};

template<class K, class V>
class AVLTree
{
	typedef AVLTreeNode<K, V> Node;

public:
    bool Insert(const pair<K, V>& kv)
    {
        if (_root == nullptr)
        {
            _root = new Node(kv);
            return true;
        }
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur)
        {
            if (cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        } 
        cur = new Node(kv);
        if (parent->_kv.first < kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }
        cur->_parent = parent;
        // 更新平衡因⼦
            while (parent)
            {
                // 更新平衡因⼦
                if (cur == parent->_left)
                {
                    parent->_bf--;
                }
                else
                {
                    parent->_bf++;
                }
                if (parent->_bf == 0)
                {
                    //  更新结束
                        break;
                }
                else if (parent->_bf == 1 || parent->_bf == -1)
                {
                    // 继续往上更新
                        cur = parent;
                    parent = parent->_parent;
                }
                else if (parent->_bf == 2 || parent->_bf == -2)
                {
                    //  不平衡了，旋转处理
                    if (parent->_bf == -2 && cur->_bf == -1)
                    {
                        RotateR(parent);
                    }
                    else if (parent->_bf == 2 && cur->_bf == 1)
                    {
                        RotateL(parent);
                    }
                    else if (parent->_bf == -2 && cur->_bf == 1)
                    {
                        RotateLR(parent);
                    }
                    else if (parent->_bf == 2 && cur->_bf == -1)
                    {
                        RotateRL(parent);
                    }
                    else
                    {
                        assert(false);
                    }

                    break;
                }
                else
                {
                    assert(false);
                }
            }
        return true;
    }

    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;//
        Node* subLR = subL->_right;//

        parent->_left = subLR;
        if (subLR)
        {
            subLR->_parent = parent;
        }

        Node* ppnode = parent->_parent;

        subL->_right = parent;
        parent->_parent = subL;

        if (parent == _root)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            if (ppnode->_left == parent)
            {
                ppnode->_left = subL;
            }
            else
            {
                ppnode->_right = subL;
            }
            subL->_parent = ppnode;
        }
        parent->_bf = 0;
        subL->_bf = 0;
    }

    
        void RotateL(Node * parent)
        {
            Node* subR = parent->_right;
            Node* subRL = subR->_left;
            parent->_right = subRL;
            if (subRL)
                subRL->_parent = parent;
            Node* parentParent = parent->_parent;
            subR->_left = parent;
            parent->_parent = subR;

            if (parentParent == nullptr)
            {
                _root = subR;
                subR->_parent = nullptr;
            }
            else
            {
                if (parent == parentParent->_left)
                {
                    parentParent->_left = subR;
                }
                else
                {
                    parentParent->_right = subR;
                }
                subR->_parent = parentParent;
            }
            parent->_bf = subR->_bf = 0;
        }

    void RotateLR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        int bf = subLR->_bf;

        RotateL(parent->_left);
        RotateR(parent);

        //平衡因子的更新
        parent->_bf = 0;
        subL->_bf = 0;

        if (bf == 1)
        {
            parent->_bf = -1;
        }
        else if (bf == -1)
        {
            subL->_bf = 1;
        }
        subLR->_bf = 0;
    }
    
    void RotateRL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        int bf = subRL->_bf;
        RotateR(parent->_right);
        RotateL(parent);

        parent->_bf = 0;
        subR->_bf = 0;

        if (bf == 1)
        {
            subR->_bf = -1;
        }
        else if (bf == -1)
        {
            parent->_bf = 1;
        }
        subRL->_bf = 0;
    }

    void InOrder()//为什么要写在public中，因为这个函数可以在类的外部使用
    {
        _InOrder(_root);
        cout << endl;
    }
private:
    void _InOrder(Node* root)
    {
        if (root == nullptr)
        {
            return;
        }
        _InOrder(root->_left);
        cout << root->_kv.first << " ";
        _InOrder(root->_right);
    }

private:
	Node* _root = nullptr;
};

void TestAVLTree1()
{
    AVLTree<int, int> t;
    // 常规的测试⽤例

        //int a[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
        // 特殊的带有双旋场景的测试⽤例

        int a[] = { 4, 2, 6, 1, 3, 5, 15, 7, 16, 14 };
    for (auto e : a)
    {
        t.Insert({ e, e });
    }
    t.InOrder();
    //cout << t.IsBalanceTree() << endl;
}